"""This is an agent-based simulation inspired by Caucutt (2006).

The economy consists of a single firm...blah blah blah.

Created by Michael Selik, selik AT gatech DOT edu."""

import math
import random


class World(object):
  """The world houses everything and defines the rules."""
  def __init__(self, N):
    self.time = 0
    self.firm = Firm(self)
    self.firm.workers = [Worker(self.firm) for n in range(N)]
    self.tutor = Tutor()

  def run(self, T):
    return [self.step() for t in range(T)]

  def step(self):
    self.time += 1
    return self.firm.produce()


class Firm(object):
  """The firm hires workers and produces goods."""
  def __init__(self, world):
    self.world = world
    self.workers = []
    self.retirees = []
    self.skilled_wage = 1
    self.unskilled_wage = 1
    self.tutor = Tutor()
    self.learner = Learner(self)

    # production constants
    self.alpha = 0.5
    self.vee = 0.5
    self.theta = 0.8

    # production variables
    self.technology = 100
    self.capital = 100

  @property
  def skilled_workers(self):
    return filter(lambda x: x.skilled, self.workers)

  @property
  def unskilled_workers(self):
    return filter(lambda x: not x.skilled, self.workers)

  @property
  def skilled_retirees(self):
    return filter(lambda x: x.skilled, self.retirees)

  @property
  def unskilled_retirees(self):
    return filter(lambda x: not x.skilled, self.retirees)

  @property
  def proportion_skilled(self):
    return len(self.skilled_workers) / len(self.workers)

  def produce(self):
    self.update_wages()
    for skilled in self.skilled_workers:
      skilled.earn(self.skilled_wage)
    for unskilled in self.unskilled_workers:
      unskilled.earn(self.unskilled_wage)
    for retiree in [worker.retire() for worker in self.workers]:
      retiree.raise_child()

  def update_wages(self):
    self.skilled_wage = self.marginal_product_labor(True)
    self.unskilled_wage = self.marginal_product_labor(False)
  
  def marginal_product_labor(self, skilled):
    Ns = len(self.skilled_workers)
    Nu = len(self.unskilled_workers)
    base = self.quantity(Ns, Nu)
    if skilled:
      return base - self.quantity(Ns - 1, Nu)
    else:
      return base - self.quantity(Ns, Nu - 1)

  def quantity(self, Ns, Nu):
    if Ns < 0:
      Ns = 0
    if Nu < 0:
      Nu = 0
    Ls = self.theta * pow(Ns, self.vee)
    Lu = (1 - self.theta) * pow(Nu, self.vee)
    L = pow(Ls + Lu, (1 - self.alpha) / self.vee)
    A = pow(self.technology, 1 - self.alpha)
    K = pow(self.capital, self.alpha)
    return A * K * L


class Tutor(object):
  """The tutor helps workers become skilled."""
  def __init__(self):
    self.difficulty = 3

  def teach(self, pupil):
    return random.random() < self.chance(pupil)

  def chance(self, pupil):
    return 1 - math.exp(-1 * self.difficulty * self.helper(pupil))

  def helper(self, pupil):
    # return pupil.parent.investment_rate
    if pupil.firm.skilled_wage > 0:
      return pupil.parent.investment / pupil.firm.skilled_wage
    else: # if there are no skilled people, no one can teach
      return 0

class Worker(object):
  """The behavior of workers is the primary interest of this model."""
  def __init__(self, firm, parent = None):
    self.firm = firm
    self.world = firm.world
    self.parent = parent
    self.child = None
    self.consumption = 0
    self.investment = 0
    self.beta = 0.9 # time discount

    if parent:
      self.skilled = self.firm.tutor.teach(self)
      self.investment_rate = self.choose_investment_rate()
    else:
      self.investment_rate = random.random()
      self.skilled = random.choice([True, False])
  
  
  @property
  def utility(self):
    if self.child:
      return self.consumption + self.beta * self.child.consumption
    else:
      return self.consumption

  def earn(self, wage):
    if self.investment_rate < 0:
      self.investment_rate = 0
    if self.investment_rate > 1:
      self.investment_rate = 1
    self.consumption += (1 - self.investment_rate) * wage
    self.investment += self.investment_rate * wage

    self.proportion_skilled = self.firm.proportion_skilled
    return wage

  def retire(self):
    self.firm.retirees.append(self)
    self.firm.workers.remove(self)
    if len(self.firm.retirees) > 6 * len(self.firm.workers):
      del self.firm.retirees[0]
    return self

  def raise_child(self):
    self.child = Worker(self.firm, self)
    self.firm.workers.append(self.child)

  def choose_investment_rate(self):
    r = self.firm.learner.best_rate(self)
    if r > 1:
      return 1
    if r < 0:
      return 0
    return r

class Learner(object):
  """The learner will figure out the best investment rate for a worker."""
  def __init__(self, firm):
    self.firm = firm
    self.best_rate_skilled = random.random()
    self.best_rate_unskilled = random.random()
    self.exploration_range = 0.25
    self.cooling_index = 0

  def best_rate(self, worker):
    if random.random() > 0.95:
      self.update_best_rate()
      self.update_exploration_range()
    if worker.skilled:
      return self.best_rate_skilled + self.exploration()
    else:
      return self.best_rate_unskilled + self.exploration()

  def update_best_rate(self):
    best_skilled = self.pick_best(self.firm.skilled_retirees)
    best_unskilled = self.pick_best(self.firm.unskilled_retirees)
    self.best_rate_skilled = best_skilled
    if self.best_rate_skilled < 0:
      self.best_rate_skilled = 0
    elif self.best_rate_skilled > 1:
      self.best_rate_skilled = 1
    self.best_rate_unskilled = best_unskilled
    if self.best_rate_unskilled < 0:
      self.best_rate_unskilled = 0
    elif self.best_rate_unskilled > 1:
      self.best_rate_unskilled = 1

  def pick_best(self, group):
    grp = filter(lambda x: \
      x.proportion_skilled + 0.1 > self.firm.proportion_skilled \
      and x.proportion_skilled - 0.1 < self.firm.proportion_skilled, group)
    if len(grp) == 0:
      return random.random()
    return max(grp, key=lambda x: x.utility).investment_rate

  def exploration(self):
    bot = self.exploration_range * -1
    top = self.exploration_range
    return random.uniform(bot, top)

  def update_exploration_range(self):
    self.cooling_index += 10
    x = self.cooling_index
    z = len(self.firm.workers) + 1
    range = math.exp(-x / z - 0.7)
    self.exploration_range = range
